/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file conv2d_intf.h
 * \brief
 */

#ifndef CONV2D_INTF_H
#define CONV2D_INTF_H

#include "kernel_utils.h"
#include "conv2d_framework_util.h"
#include "conv2d_common_func.h"

namespace conv2d {

template<class Config, template<typename, class>class Impl>
struct Conv2dIntf {
    using Ext = Impl<Conv2dIntf, Config>;
    using FmapT = typename Config::FmapT;
    using WeightT = typename Config::WeightT;
    using OutputT = typename Config::OutputT;
    using BiasT = typename Config::BiasT;
    using ScaleT = typename Config::ScaleT;
    using L0cT = typename Config::L0cT;
    using ContextType = typename Ext::ContextData;
    using ImplDataType = typename Ext::ImplDataType;
    using ConvParam = typename Config::ConvParam;
public:
    ContextType ctx;
    ImplDataType impl;

    __aicore__ inline Conv2dIntf() {}

    __aicore__ inline void Init(const void* __restrict conv2dTiling)
    {
        using local = typename Ext::Init;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, conv2dTiling)) {
            local::call(this, conv2dTiling);
        }
    }

    __aicore__ inline void SetFmap(const GlobalTensor<FmapT>& fmap)
    {
        using local = typename Ext::SetFmap;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, fmap)) {
            local::call(this, fmap);
        }
    }

    __aicore__ inline void SetWeight(const GlobalTensor<WeightT>& weight)
    {
        using local = typename Ext::SetWeight;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, weight)) {
            local::call(this, weight);
        }
    }

    __aicore__ inline void SetScale(const GlobalTensor<ScaleT>& scale)
    {
        using local = typename Ext::SetScale;
        if constexpr(CONV_CHECK_FUN(local, Conv2dFunc, this, scale)) {
            local::call(this, scale);
        }
    }

    __aicore__ inline void SetBias(const GlobalTensor<BiasT>& bias)
    {
        using local = typename Ext::SetBias;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, bias)) {
            local::call(this, bias);
        }
    }

    __aicore__ inline void SetOrgFmapShape(uint64_t orgCi, uint64_t orgHi, uint64_t orgWi)
    {
        using local = typename Ext::SetOrgFmapShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, orgCi, orgHi, orgWi)) {
            local::call(this, orgCi, orgHi, orgWi);
        }
    }

    __aicore__ inline void SetOrgWeightShape(uint64_t orgCo, uint64_t orgCi, uint64_t orgKh, uint64_t orgKw)
    {
        using local = typename Ext::SetOrgWeightShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, orgCo, orgCi, orgKh, orgKw)) {
            local::call(this, orgCo, orgCi, orgKh, orgKw);
        }
    }

    __aicore__ inline void SetOrgOutputShape(uint64_t orgCo, uint64_t orgHo, uint64_t orgWo)
    {
        using local = typename Ext::SetOrgOutputShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, orgCo, orgHo, orgWo)) {
            local::call(this, orgCo, orgHo, orgWo);
        }
    }

    __aicore__ inline void SetSingleFmapShape(uint64_t singleCi, uint64_t singleHi, uint64_t singleWi)
    {
        using local = typename Ext::SetSingleFmapShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, singleCi, singleHi, singleWi)) {
            local::call(this, singleCi, singleHi, singleWi);
        }
    }

    __aicore__ inline void SetSingleWeightShape(uint64_t singleCo, uint64_t singleCi, uint64_t singleKh,
                                                uint64_t singleKw)
    {
        using local = typename Ext::SetSingleWeightShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, singleCo, singleCi, singleKh, singleKw)) {
            local::call(this, singleCo, singleCi, singleKh, singleKw);
        }
    }

    __aicore__ inline void SetSingleOutputShape(uint64_t singleCo, uint64_t singleHo, uint64_t singleWo)
    {
        using local = typename Ext::SetSingleOutputShape;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, singleCo, singleHo, singleWo)) {
            local::call(this, singleCo, singleHo, singleWo);
        }
    }

    __aicore__ inline void SetFmapStartPosition(int64_t singleCoreHiStartPos, int64_t wiStartPos)
    {
        using local = typename Ext::SetFmapStartPosition;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this, singleCoreHiStartPos, wiStartPos)) {
            local::call(this, singleCoreHiStartPos, wiStartPos);
        }
    }

    template <bool sync = true>
    __aicore__ inline void IterateAll(const GlobalTensor<OutputT>& output, bool enPartialSum = false)
    {
        using local = typename Ext::IterateAll;
        if constexpr (CONV_CHECK_FUN_TEMPLATE(local, Conv2dFunc, sync, this, output, enPartialSum)) {
            local::template call<sync>(this, output, enPartialSum);
        }
    }

    __aicore__ inline void End()
    {
        using local = typename Ext::End;
        if constexpr (CONV_CHECK_FUN(local, Conv2dFunc, this)) {
            local::call(this);
        }
    }

private:
    template <bool sync = true>
    __aicore__ inline bool Iterate(bool enPartialSum = false)
    {
        using local = typename Ext::Iterate;
        if constexpr (CONV_CHECK_FUN_TEMPLATE(local, Conv2dFunc, sync, this, enPartialSum)) {
            return local::template call<sync>(this, enPartialSum);
        }
        return false;
    }

    template <bool sync = true>
    __aicore__ inline void GetTensorC(const GlobalTensor<OutputT>& output, bool enSequentialWrite = false)
    {
        using local = typename Ext::GetTensorC;
        if constexpr (CONV_CHECK_FUN_TEMPLATE(local, Conv2dFunc, sync, this, output, enSequentialWrite)) {
            local::template call<sync>(this, output, enSequentialWrite);
        }
    }
};

}  // namespace conv2d

#endif // __CONV2D_INTF_H__
